# mcast.py
import socket
import datetime
#import netifaces
import sys
import time
import struct

IPV4_MULTICAST_ADDR = '224.0.0.120'
IPV6_MULTICAST_ADDR = "FF02::A1F7"
MULTICAST_PORT = 911

IPV4_LOCAL_ADDR = '10.25.12.111'
IPV6_LOCAL_ADDR = "fec0:10.25.12::111"
MAX_SIZE = 10240

START_ABSOLUTE_TIME = datetime.datetime.now()


# 定义出错和报告接口
def fatal_error(message):
    sys.exit(message)


def report(message):
    print("{}".format(message))


# 获取接口IPV4地址
def interface_ipv4_address(interface_name="NORMAL"):

    if interface_name == "NORMAL":
        address_str = IPV4_LOCAL_ADDR
    else:
        pass
        # interface_addresses = netifaces.interfaces()
        # if not interface_name in netifaces.interfaces():
        #     fatal_error("Interface " + interface_name + " not present.")
        # interface_addresses = netifaces.ifaddresses(interface_name)
        # if not netifaces.AF_INET in interface_addresses:
        #     fatal_error(
        #         "Interface {} has no IPv4 address.".format(interface_name))
        # address_str = interface_addresses[netifaces.AF_INET][0]['addr']
    return address_str


# 获取接口IPV6地址
def interface_ipv6_address(interface_name="NORMAL", exclude_scope=False):
    if interface_name == "NORMAL":
        address_str = IPV6_LOCAL_ADDR
    else:
        pass
        # interface_addresses = netifaces.interfaces()
        # if not interface_name in netifaces.interfaces():
        #     fatal_error("Interface " + interface_name + " not present.")
        # interface_addresses = netifaces.ifaddresses(interface_name)
        # if not netifaces.AF_INET6 in interface_addresses:
        #     fatal_error(
        #         "Interface {} has no IPv6 address.".format(interface_name))
        # address_str = interface_addresses[netifaces.AF_INET6][0]['addr']
        # if exclude_scope:
        #     if "%" in address_str:
        #         address_str = address_str.split("%")[0]
    return address_str


# 允许IP和PORT重复绑定在不同的套接字上
def enable_addr_and_port_reuse(sock):
    try:
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    except AttributeError:
        pass
    try:
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
    except AttributeError:
        pass


def create_ipv4_mcast_rx_socket(interface_name="NORMAL"):
    # pylint:disable=no-member
    report("Create IPv4 RX socket for group {} on {}".format(
        IPV4_MULTICAST_ADDR, interface_name))
    local_address = interface_ipv4_address(interface_name)
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM,
                         socket.IPPROTO_UDP)
    enable_addr_and_port_reuse(sock)
    # sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, 32)
    # sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_LOOP, 1)
    sock.bind(('', MULTICAST_PORT))

    req = struct.pack("=4s4s", socket.inet_aton(IPV4_MULTICAST_ADDR),
                      socket.inet_aton(local_address))
    sock.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, req)
    return sock


def create_ipv6_mcast_rx_socket(interface_name="NORMAL"):
    # pylint:disable=no-member
    report("Create IPv6 RX socket for group {} on {}".format(
        IPV6_MULTICAST_ADDR, interface_name))
    sock = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM,
                         socket.IPPROTO_UDP)
    enable_addr_and_port_reuse(sock)
    sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_PKTINFO, 1)
    sock.bind(("::", MULTICAST_PORT))
    #interface_index = socket.if_nametoindex(interface_name)
    mreq = struct.pack("16s15s".encode('utf-8'), 
        socket.inet_pton(socket.AF_INET6, IPV6_MULTICAST_ADDR), (chr(0) * 16).encode('utf-8'))
    sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_JOIN_GROUP, mreq)
    return sock


def create_ipv4_mcast_tx_socket(interface_name="NORMAL"):
    report("Create IPv4 TX socket for group {} on {}".format(
        IPV4_MULTICAST_ADDR, interface_name))
    local_address = interface_ipv4_address(interface_name)
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM,
                         socket.IPPROTO_UDP)
    enable_addr_and_port_reuse(sock)
    sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_IF,
                    socket.inet_aton(local_address))
    sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_LOOP, 0)
    sock.bind(("0.0.0.0", MULTICAST_PORT))
    sock.connect((IPV4_MULTICAST_ADDR, MULTICAST_PORT))
    return sock


def create_ipv6_mcast_tx_socket(interface_name="NORMAL"):
    report("Create IPv6 TX socket for group {} on {}".format(
        IPV6_MULTICAST_ADDR, interface_name))
    sock = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM,
                         socket.IPPROTO_UDP)
    enable_addr_and_port_reuse(sock)
    #interface_index = socket.if_nametoindex(interface_name)
    sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_MULTICAST_IF,(chr(0) * 16).encode('utf-8'))
    sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_MULTICAST_LOOP, 0)
    sock.bind(("::", MULTICAST_PORT))
    sock.connect((IPV6_MULTICAST_ADDR, MULTICAST_PORT))
    return sock


def send_pkt(protol, message):
    i = 0
    if protol == "ipv4":
        try:
            sock = create_ipv4_mcast_tx_socket()
        except Exception as exception:
            report("create {} socket failed".format(protol))
            #return
    elif protol == "ipv6":
        try:
            sock = create_ipv6_mcast_tx_socket()
        except Exception as exception:
            report("create {} socket failed".format(protol))
            return
    else:
        return
    while True:
        try:
            sock.send(message.encode())
            time.sleep(1)
            print("send {}".format(i))
            i = i + 1
        except Exception as exception:
            report("send {} message failed".format(protol))


def recv_pkt(protol, result):
    i = 10
    if protol == "ipv4":
        try:
            sock = create_ipv4_mcast_rx_socket()
        except Exception as exception:
            report("create {} socket failed,because {}".format(protol, exception))
            print("error")
    elif protol == "ipv6":
        try:
            sock = create_ipv6_mcast_rx_socket()
        except Exception as exception:
            report("create {} socket failed, because {}".format(protol, exception))
            print("error")
    else:
        print("error") 
    for i in range(10):
        try:
            message,addr = sock.recvfrom(MAX_SIZE)
            time.sleep(1)
        except Exception as e:
            report("create {} socket failed, because {}".format(protol, e))
        else:
            message_str = message.decode()
            # print(message_str)
            # print(result)
            if result == message_str:
                print("ok") 
            else:
                print("error") 

def print_help():
    print("\r\npython mcast.py [option] [family] [message_length] "\
              "\r\noption:" \
              "\r\n\t\t-h : print this help" \
              "\r\n\t\t-s : send multicast message"\
              "\r\n\t\t-r : receive multicast message" \
              "\r\nfamily:"
              "\r\n\t\tv4 : ipv4 multicast message" \
              "\r\n\t\tv6 : ipv6 multicast message" \
              "\r\nmessage_length: [0-65535]")

if __name__ == "__main__":
    # send_pkt("ipv6","yanghuajun")
    # recv_pkt("ipv4")
    if len(sys.argv) != 4:
        print_help()
    else:
        if sys.argv[1] == "-s" and sys.argv[2] == "v4" and 65535 > int(sys.argv[3]) > 0:
            send_pkt("ipv4","a"*int(sys.argv[3]))
        elif sys.argv[1] == "-s" and sys.argv[2] == "v6" and 65535 > int(sys.argv[3]) > 0:
            send_pkt("ipv6","a"*int(sys.argv[3]))
        elif sys.argv[1] == "-r" and sys.argv[2] == "v4" and 65535 > int(sys.argv[3]) > 0:
            recv_pkt("ipv4","a"*int(sys.argv[3]))
        elif sys.argv[1] == "-r" and sys.argv[2] == "v6" and 65535 > int(sys.argv[3]) > 0:
            recv_pkt("ipv6","a"*int(sys.argv[3]))    
        else:
            print("para error: please input \"-h\" get help")   
 
    # print()
    # for i in sys.argv:
    #     print(i)
       